import cv2
import numpy as np
import os

class MangaPanelExtractor:
    def __init__(self, min_panel_size=500):
        self.min_panel_size = min_panel_size
        self.debug = True

    def preprocess_image(self, image):
        """预处理图片，增强边缘和轮廓"""
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 使用中值滤波去除噪点
        denoised = cv2.medianBlur(gray, 3)
        
        # 使用自适应阈值处理
        binary = cv2.adaptiveThreshold(
            denoised,
            255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY,
            11,
            2
        )
        
        # 使用Canny边缘检测
        edges = cv2.Canny(binary, 50, 150)
        
        # 膨胀边缘，使轮廓更明显
        kernel = np.ones((3,3), np.uint8)
        dilated = cv2.dilate(edges, kernel, iterations=2)
        
        return dilated, binary

    def find_contours(self, edge_image, original_image):
        """查找并过滤轮廓"""
        height, width = edge_image.shape[:2]
        contours, hierarchy = cv2.findContours(
            edge_image,
            cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_SIMPLE
        )
        
        valid_contours = []
        for cnt in contours:
            # 计算轮廓面积
            area = cv2.contourArea(cnt)
            if area < self.min_panel_size:
                continue
                
            # 获取最小外接矩形
            x, y, w, h = cv2.boundingRect(cnt)
            
            # 过滤掉太小或太大的矩形
            if (w * h < self.min_panel_size or 
                w < width * 0.2 or  # 宽度至少为图片宽度的20%
                h < height * 0.1):  # 高度至少为图片高度的10%
                continue
                
            # 计算矩形区域内的平均亮度
            roi = original_image[y:y+h, x:x+w]
            avg_brightness = np.mean(roi)
            
            # 如果是纯白色区域，跳过
            if avg_brightness > 250:
                continue
                
            valid_contours.append((x, y, w, h))
        
        return valid_contours

    def merge_overlapping_panels(self, panels):
        """合并重叠的分镜"""
        if not panels:
            return []
            
        # 按y坐标排序
        panels = sorted(panels, key=lambda x: x[1])
        
        merged = []
        current = list(panels[0])
        
        for panel in panels[1:]:
            x, y, w, h = panel
            curr_x, curr_y, curr_w, curr_h = current
            
            # 计算当前矩形的底部y坐标
            curr_bottom = curr_y + curr_h
            
            # 如果当前矩形与下一个矩形重叠
            if y < curr_bottom:
                # 更新当前矩形的宽度和高度
                current[2] = max(curr_w, w)
                current[3] = max(curr_bottom, y + h) - curr_y
            else:
                # 如果不重叠，保存当前矩形并开始新的矩形
                merged.append(tuple(current))
                current = list(panel)
        
        # 添加最后一个矩形
        merged.append(tuple(current))
        return merged

    def extract_panels(self, image_path, output_dir):
        """提取漫画分镜"""
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        print(f"\n开始处理图片: {image_path}")
        
        # 读取图片
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图片: {image_path}")
            
        # 获取图片尺寸
        height, width = image.shape[:2]
        print(f"图片尺寸: {width}x{height}")
        
        # 预处理图片
        print("预处理图片...")
        edge_image, binary = self.preprocess_image(image)
        
        # 保存预处理结果用于调试
        if self.debug:
            cv2.imwrite(os.path.join(output_dir, "edges.jpg"), edge_image)
            cv2.imwrite(os.path.join(output_dir, "binary.jpg"), binary)
        
        # 查找轮廓
        print("查找分镜轮廓...")
        panels = self.find_contours(edge_image, binary)
        print(f"找到 {len(panels)} 个初始分镜")
        
        # 合并重叠的分镜
        panels = self.merge_overlapping_panels(panels)
        print(f"合并后剩余 {len(panels)} 个分镜")
        
        # 在原图上绘制检测到的分镜位置
        debug_image = image.copy()
        for i, (x, y, w, h) in enumerate(panels):
            cv2.rectangle(debug_image, (x, y), (x+w, y+h), (0, 255, 0), 2)
            cv2.putText(debug_image, str(i+1), (x+10, y+30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        # 保存调试图片
        cv2.imwrite(os.path.join(output_dir, "detected_panels.jpg"), debug_image)
        
        # 提取并保存每个分镜
        print("\n保存分镜...")
        for i, (x, y, w, h) in enumerate(panels):
            # 提取分镜
            panel = image[y:y+h, x:x+w]
            
            # 保存分镜
            output_path = os.path.join(output_dir, f"panel_{i+1}.jpg")
            cv2.imwrite(output_path, panel)
            print(f"保存分镜 {i+1}: {output_path}")
        
        return len(panels)

def main():
    # 创建实例
    extractor = MangaPanelExtractor(min_panel_size=5000)
    
    # 确保输入和输出目录存在
    input_dir = "pic"
    output_base_dir = "extracted_panels"
    
    if not os.path.exists(input_dir):
        os.makedirs(input_dir)
        print(f"创建输入目录: {input_dir}")
    
    if not os.path.exists(output_base_dir):
        os.makedirs(output_base_dir)
        print(f"创建输出目录: {output_base_dir}")
    
    # 获取所有PNG文件
    png_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.jpg')]
    print(f"\n找到 {len(png_files)} 个PNG文件")
    
    if not png_files:
        print(f"\n错误：在 {input_dir} 目录中没有找到PNG文件！")
        print("请确保：")
        print("1. 图片文件已经放在 pic 目录中")
        print("2. 图片文件是PNG格式的")
        print("3. 文件扩展名是.png（小写）")
        return
    
    # 处理每个PNG文件
    for png_file in png_files:
        print(f"\n处理文件: {png_file}")
        input_path = os.path.join(input_dir, png_file)
        
        # 创建输出目录
        output_dir = os.path.join(output_base_dir, os.path.splitext(png_file)[0])
        
        try:
            # 如果输出目录已存在，先删除
            if os.path.exists(output_dir):
                import shutil
                shutil.rmtree(output_dir)
                print(f"清理已存在的输出目录: {output_dir}")
            
            # 处理图片
            num_panels = extractor.extract_panels(input_path, output_dir)
            print(f"成功从 {png_file} 中提取了 {num_panels} 个分镜到 {output_dir}")
            
        except Exception as e:
            print(f"处理 {png_file} 时出错: {str(e)}")
            continue

if __name__ == "__main__":
    main() 